import numpy as np
import weakref
from functools import reduce
from typing import Any
from typing import List as TypingList
from typing import Optional

_epsilon = 1e-6


def get_bounds_fit_matrix(
        xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float
) -> np.ndarray:
    """Return a 4x4 transform matrix mapping the default bounding box [-0.5, 0.5, ...] into a custom bounding box.

    Parameters
    ----------
    xmin, xmax, ymin, ymax, zmin, zmax : float
        Bounds for the target box.

    Returns
    -------
    ndarray
        4x4 transformation matrix.
    """
    # Validate all arguments are floats
    for name, value in locals().copy().items():
        try:
            float(value)
        except (TypeError, ValueError):
            raise TypeError(
                "%s: expected float, %s given" % (name, type(value).__name__)
            )
    # Create diagonal scaling matrix and set translation
    matrix = np.diagflat(
        np.array((xmax - xmin, ymax - ymin, zmax - zmin, 1.0), np.float32, order="C")
    )
    matrix[0:3, 3] = ((xmax + xmin) / 2.0, (ymax + ymin) / 2.0, (zmax + zmin) / 2.0)
    return matrix


class Transform(object):
    """
    Abstraction of a 4x4 model transformation matrix with hierarchy support.

    Supports translation, rotation (as quaternion), scaling, custom matrix, and parent-child relationships.
    """

    def __init__(
            self,
            bounds: Optional[TypingList[float]] = None,
            translation: Optional[TypingList[float]] = None,
            rotation: Optional[TypingList[float]] = None,
            scaling: Optional[TypingList[float]] = None,
            custom_matrix: Optional[np.ndarray] = None,
            parent: Optional["Transform"] = None,
    ):
        """
        Initialize a Transform object.

        Parameters
        ----------
        bounds : list of float, optional
            Bounding box as [xmin, xmax, ymin, ymax, zmin, zmax].
        translation : list of float, optional
            Translation vector [dx, dy, dz].
        rotation : list of float, optional
            Quaternion as [gamma, axis_x, axis_y, axis_z].
        scaling : list of float, optional
            Scaling coefficients [sx, sy, sz].
        custom_matrix : np.ndarray, optional
            4x4 custom transformation matrix.
        parent : Transform, optional
            Parent transform for hierarchy.
        """
        # Validate and assign parameters
        self.bounds = bounds
        assert translation is None or len(translation) == 3
        self.translation = translation
        assert rotation is None or len(rotation) == 4
        self.rotation = rotation
        assert scaling is None or len(scaling) == 3
        self.scaling = scaling
        self.parent = parent
        if parent is not None:
            parent._add_child(self)
        self.drawables: TypingList[weakref.ref] = []
        self.children: TypingList[weakref.ref] = []
        self.parent_matrix = (
            parent.model_matrix if parent else np.identity(4, dtype=np.float32)
        )
        self.custom_matrix = (
            custom_matrix
            if custom_matrix is not None
            else np.identity(4, dtype=np.float32)
        )
        self.model_matrix = np.identity(4, dtype=np.float32)
        self._recompute_matrix()

    def __setattr__(self, key: str, value: Any) -> None:
        """Set attributes with conversion to ndarray where needed. Triggers matrix recomputation on update."""
        is_set = hasattr(self, key)  # == False in constructor
        # parameter canonicalization and some validation via reshaping
        if value is None:
            # Forbid None for critical transform fields that should always have valid values
            if key in ["parent_matrix", "custom_matrix", "model_matrix"]:
                raise ValueError(
                    f"Cannot set {key} to None. These fields must have valid matrix values."
                )
            # Allow None for optional transform parameters (translation, rotation, scaling, bounds)
        elif key == "translation":
            # Ensure translation is a 3x1 column vector
            value = np.array(value, dtype=np.float32).reshape(3, 1)
        elif key == "rotation":
            # Convert rotation to quaternion, normalize, and ensure valid axis
            value = np.array(value, dtype=np.float32).reshape(4)
            value[0] = np.fmod(value[0], 2.0 * np.pi)
            if value[0] < 0.0:
                value[0] += 2.0 * np.pi
            value[0] = np.cos(value[0] / 2)
            norm = np.linalg.norm(value[1:4])
            needed_norm = np.sqrt(1 - value[0] * value[0])
            if abs(norm - needed_norm) > _epsilon:
                if norm < _epsilon:
                    raise ValueError(
                        "Norm of (x, y, z) part of quaternion too close to zero"
                    )
                value[1:4] = value[1:4] / norm * needed_norm
            # assert abs(np.linalg.norm(value) - 1.0) < _epsilon
        elif key == "scaling":
            # Ensure scaling is a 3-element vector
            value = np.array(value, dtype=np.float32).reshape(3)
        elif key in ["parent_matrix", "custom_matrix", "model_matrix"]:
            # Ensure all matrices are 4x4
            value = np.array(value, dtype=np.float32).reshape((4, 4))
        super(Transform, self).__setattr__(key, value)
        if is_set and key != "model_matrix":
            self._recompute_matrix()
            self._notify_dependants()

    def __repr__(self) -> str:
        return "Transform(bounds={!r}, translation={!r}, rotation={!r}, scaling={!r})".format(
            self.bounds, self.translation, self.rotation, self.scaling
        )

    def _recompute_matrix(self) -> None:
        """Recompute the model matrix from all transform parameters."""
        # this method shouldn't modify any fields except self.model_matrix
        if self.bounds is None or len(self.bounds) == 0:
            fit_matrix = np.identity(4)
        else:
            if len(self.bounds) == 6:
                xmin, xmax, ymin, ymax, zmin, zmax = self.bounds
            elif len(self.bounds) == 4:
                xmin, xmax, ymin, ymax = self.bounds
                zmin, zmax = -0.5, 0.5
            elif len(self.bounds) == 2:
                # 1-D bounds for a 2D strip - why not?
                xmin, xmax = self.bounds
                ymin, ymax, zmin, zmax = -0.5, 0.5, -0.5, 0.5
            else:
                raise ValueError(
                    "Wrong size of bounds array ({}), should be 4 for 2D or 6 for 3D bounds.".format(
                        self.bounds
                    )
                )
            fit_matrix = get_bounds_fit_matrix(xmin, xmax, ymin, ymax, zmin, zmax)
        # Build translation matrix
        if self.translation is not None:
            translation_matrix = np.vstack(
                (
                    np.hstack(
                        (np.identity(3), np.array(self.translation).reshape(3, 1))
                    ),
                    np.array([0.0, 0.0, 0.0, 1.0]).reshape(1, 4),
                )
            )
        else:
            translation_matrix = np.identity(4)
        # Build rotation matrix from quaternion
        if self.rotation is not None:
            a, b, c, d = self.rotation
            rotation_matrix = np.array(
                [
                    [
                        a * a + b * b - c * c - d * d,
                        2 * (b * c - a * d),
                        2 * (b * d + a * c),
                        0.0,
                    ],
                    [
                        2 * (b * c + a * d),
                        a * a - b * b + c * c - d * d,
                        2 * (c * d - a * b),
                        0.0,
                    ],
                    [
                        2 * (b * d - a * c),
                        2 * (c * d + a * b),
                        a * a - b * b - c * c + d * d,
                        0.0,
                    ],
                    [0.0, 0.0, 0.0, 1.0],
                ]
            )
        else:
            rotation_matrix = np.identity(4)
        # Build scaling matrix
        if self.scaling is not None:
            scaling_matrix = np.diag(np.append(self.scaling, 1.0))
        else:
            scaling_matrix = np.identity(4)
        # Compose all matrices in the correct order
        self.model_matrix = reduce(
            np.dot,
            [
                translation_matrix,
                rotation_matrix,
                scaling_matrix,
                fit_matrix,
                self.custom_matrix,
                self.parent_matrix,
            ],
        )

    def _add_child(self, transform: "Transform") -> None:
        self.children.append(weakref.ref(transform))

    def add_drawable(self, drawable: Any) -> None:
        """Register a Drawable to have its model_matrix overwritten after changes to the transform or its parent."""
        self.drawables.append(weakref.ref(drawable))

    def parent_updated(self) -> None:
        """Read updated parent transform matrix and update own model_matrix.

        This method should be normally only called by parent Transform to notify its children.
        """
        if self.parent is not None:
            self.parent_matrix = self.parent.model_matrix
        self._recompute_matrix()
        self._notify_dependants()

    def _notify_dependants(self) -> None:
        extant_children = []
        for child_ref in self.children:
            child = child_ref()
            if child is not None:
                child.parent_updated()
                extant_children.append(child_ref)
        self.children[:] = extant_children

        extant_drawables = []
        for drawable_ref in self.drawables:
            drawable = drawable_ref()
            if drawable is not None:
                drawable.model_matrix = self.model_matrix
                extant_drawables.append(drawable_ref)
        self.drawables[:] = extant_drawables


def process_transform_arguments(drawable: Any, **kwargs: Any) -> Any:
    """Create a Transform to apply to a drawable.

    Parameters
    ----------
    drawable : object
        Drawable object.
    **kwargs
        transform : Transform
            Existing Trasnform to be (re-)used on the drawable.
        xmin : float
            Lower bound in x dimsension.
        xmax : float
            Upper bound in x dimsension.
        ymin : float
            Lower bound in y dimsension.
        ymax : float
            Upper bound in y dimsension.
        zmin : float
            Lower bound in z dimsension.
        zmax : float
            Upper bound in z dimsension.
        bounds: array_like
            Dimensions bounds taking precedence over separate bound arguments,
            [xmin, xmax, ymin, ymax, zmin, zmax] or [xmin, xmax, ymin, ymax].
        translation : array_like
            [tx, ty, tz] translation vector.
        rotation : array_like
            [gamme, rx, ry, rz] rotation vector (radians).
        scaling : array_like
            [sx, sy, sz] scaling coefficients.
        model_matrix : ndarray
            Matrix of numbers.
            Must have four columns.

    Returns
    -------
    Drawable
        Transformed Drawable.

    Raises
    ------
    ValueError
        Provided transform argument is not a Transform object.
    """
    if "transform" in kwargs:
        transform = kwargs["transform"]
        if not isinstance(transform, Transform):
            raise ValueError("Provided transform argument is not a Transform object")
    else:
        separate_bounds = [
            kwargs.get("xmin", -0.5),
            kwargs.get("xmax", 0.5),
            kwargs.get("ymin", -0.5),
            kwargs.get("ymax", 0.5),
            kwargs.get("zmin", -0.5),
            kwargs.get("zmax", 0.5),
        ]

        transform_kwargs = dict(
            bounds=kwargs.get("bounds", separate_bounds),
            translation=kwargs.get("translation"),
            rotation=kwargs.get("rotation"),
            scaling=kwargs.get("scaling"),
            custom_matrix=kwargs.get("model_matrix"),
        )
        transform = Transform(**transform_kwargs)

    transform.add_drawable(drawable)
    transform.parent_updated()
    drawable.transform = transform

    return drawable


def transform(
        bounds: Optional[TypingList[float]] = None,
        translation: Optional[TypingList[float]] = None,
        rotation: Optional[TypingList[float]] = None,
        scaling: Optional[TypingList[float]] = None,
        custom_matrix: Optional[np.ndarray] = None,
        parent: Optional[Transform] = None,
) -> Transform:
    """Return a Transform object.

    Parameters
    ----------
    bounds: array_like, optional
            Dimensions bounds taking precedence over separate bound arguments, by default None.
            [xmin, xmax, ymin, ymax, zmin, zmax] or [xmin, xmax, ymin, ymax].
    translation : array_like, optional
        [tx, ty, tz] translation vector, by default None.
    rotation : array_like, optional
        [gamme, rx, ry, rz] rotation vector (radians), by default None.
    scaling : array_like, optional
        [sx, sy, sz] scaling coefficients, by default None.
    custom_matrix : _type_, optional
        Matrix of numbers, by default None.
        Must have four columns.
    parent : Transform, optional
        Optional parent transform, by default None.

    Returns
    -------
    Transform
        Transform object.
    """
    return Transform(bounds, translation, rotation, scaling, custom_matrix, parent)
